/*
 *
 *
Copyright (c) 1998-2001 Aristar, Inc. All rights reserved.

Redistribution and use in source and binary forms, with or without modification,
are permitted provided that the following conditions are met:

1. Redistributions of source code must retain the above copyright notice,
this list of conditions and the following disclaimer.

2. Redistributions in binary form must reproduce the above copyright notice,
this list of conditions and the following disclaimer in the documentation
and/or other materials provided with the distribution.

3. The end-user documentation included with the redistribution, if any,
must include the following acknowledgment:

"This product includes software developed by Aristar, Inc. (http://www.aristar.com/)."

Alternately, this acknowledgment may appear in the software itself, if and
wherever such third-party acknowledgments normally appear.

4. The names "Dialect" and "Aristar, Inc." must not be used to endorse or
promote products derived from this software without prior written permission.
For written permission, please contact info@aristar.com.

5. Products derived from this software may not be called "Dialect", nor may
"Dialect" appear in their name, without prior written permission of Aristar, Inc.

THIS SOFTWARE IS PROVIDED "AS IS" AND ANY EXPRESSED OR IMPLIED WARRANTIES,
INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY
AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL
ARISTAR, INC. OR ITS CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE,
EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

This software consists of voluntary contributions made by many individuals
on behalf of Aristar, Inc.  For more information on Aristar, Inc.,
please see <http://www.aristar.com/>.

Portions of this software are based upon software freely available and
written by individuals not associated with Aristar, Inc.  Copyrights are
maintained in the relevant source code files.
 *
 *
 */

#include "stdafx.h"
#include "DConfigGlobals.h"

#ifdef CRYPTO_SUPPORT

#include "CryptoModule.h"
#include "Interpreter.h"

///////////////////////////////////
#if defined(_DEBUG) && (defined(__WINCE__) || defined(_WINNT__))
#define new DEBUG_NEW
#undef THIS_FILE
static char THIS_FILE[] = __FILE__;
#endif
///////////////////////////////////

#define XERR(x) pInterp->RuntimeError(_T(x))
#define ERR(x) {XERR(x);return AllocNil();}
#define TERR(x) {pInterp->RuntimeError(x);return AllocNil();}

static LPCTSTR lpzWrongNumArgs = _T("Wrong number of arguments to function");
static LPCTSTR lpzExpectString = _T("Expected string argument to function.");
static LPCTSTR lpzExpectInt = _T("Expected integer argument to function");
static LPCTSTR lpzExpectArray = _T("Expected array argument to function");
static LPCTSTR lpzExpectNum = _T("Expected number argument to function");

#define REQUIRE_ARGS(x)		if (argCount != x) TERR(lpzWrongNumArgs)
#define REQUIRE_STRING(x)	if (args[x].Type() != stString) TERR(lpzExpectString)
#define REQUIRE_INTEGER(x)	if ((args[x].Type() != stInteger)) TERR(lpzExpectInt)
#define REQUIRE_ARRAY(x)	if ((args[x].Type() != stArray)) TERR(lpzExpectArray)
#define REQUIRE_NUMBER(x)	if ((args[x].Type() != stInteger) && (args[x].Type() != stFloat)) TERR(lpzExpectNum)

/*#define XERR(x) pInterp->RuntimeError(_T(x))
#define ERR(x) {XERR(x);return AllocNil();}

#define REQUIRE_ARGS(x) if (argCount!=x) ERR("Wrong number of arguments to function.")
#define REQUIRE_STRING(x) if (args[x].Type() != stString) ERR("Expected string in argument(s) to function.")
#define REQUIRE_FLOAT(x) if ((args[x].Type() != stFloat)||(args[x].Type() != stInteger)) ERR("Expected number in argument(s) to function.")
#define REQUIRE_INTEGER(x) if ((args[x].Type() != stInteger)) ERR("Expected integer in arguments(s) to function.")
#define REQUIRE_ARRAY(x) if ((args[x].Type() != stArray)) ERR("Expected array in arguments(s) to function.")
*/

static VALUE vMD5ClassBase;
static VALUE vSHAClassBase;
static VALUE vBlowfishClassBase;
static VALUE vDiffieHellmanClassBase;

/////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////

#include "Blowfish.h"
#include "CImplementer.h" // For chunk encoding types

static void freeBlowfish(void* pVoid)
{
	if (pVoid) delete (CBlowfish*)pVoid;
}

CORE_FUNC(Blowfish__Init__)
{
	REQUIRE_ARGS(1);
	ASSERT_SELF(CRYPTO_MODULE_NAME);

	CBlowfish* p = NULL;
	if (args[0].Type() == stString) {
		VALUE temp = AllocChunk(args[0].String(), enc_hex, it_ui1);
		p = new CBlowfish(temp.GetChunkData(), temp.GetChunkByteCount());
	}
	else if (args[0].Type() == stChunk) {
		p = new CBlowfish(args[0].GetChunkData(), args[0].GetChunkByteCount());
	}
	else {
		ERR("Argument to blowfish constructor must be hex string or chunk");
	}
	self->Set(_T("#pBlowfish"), AllocVoid(p, freeBlowfish, _T("blowfish")));
	return AllocNil();
}

CORE_FUNC(Blowfish_Encrypt)
{
	REQUIRE_ARGS(1);
	ASSERT_SELF(CRYPTO_MODULE_NAME);

	CBlowfish* p = (CBlowfish*)self->Get(_T("#pBlowfish")).Void();
	if (p == NULL) {
		ERR("Blowfish class has not been instantiated.");
	}

	VALUE vInput = args[0];
	if (vInput.Type() == stString) {
#ifdef UNDER_CE
		CString text = args[0].String();
		int length = text.GetLength();
		BYTE *buf = new BYTE[length];
		for (int idx = 0; idx < length; ++idx) {
			buf[idx] = (BYTE)(text[idx]);
		}
		vInput = AllocChunk(length, buf);
		delete[] buf;
#else
		CString text = args[0].String();
		vInput = AllocChunk(text.GetLength(), (U8*)(LPCTSTR)text);
#endif
	}
	else if (vInput.Type() != stChunk) {
		ERR("Arg to encrypt must be chunk or string");
	}

	VALUE vOutput = vInput.Clone();
	p->Encrypt((unsigned long*)vOutput.GetChunkData(), vOutput.GetChunkByteCount() >> 3);
	return vOutput;
}

CORE_FUNC(Blowfish_Decrypt)
{
	REQUIRE_ARGS(1);
	ASSERT_SELF(CRYPTO_MODULE_NAME);

	CBlowfish* p = (CBlowfish*)self->Get(_T("#pBlowfish")).Void();
	if (p == NULL) {
		ERR("Blowfish class has not been instantiated.");
	}

	VALUE vInput = args[0];
	if (vInput.Type() == stString) {
#ifdef UNDER_CE
		CString text = args[0].String();
		int length = text.GetLength();
		BYTE *buf = new BYTE[length];
		for (int idx = 0; idx < length; ++idx) {
			buf[idx] = (BYTE)(text[idx]);
		}
		vInput = AllocChunk(length, buf);
		delete[] buf;
#else
		CString text = args[0].String();
		vInput = AllocChunk(text.GetLength(), (U8*)(LPCTSTR)text);
#endif
	}
	else if (vInput.Type() != stChunk) {
		ERR("Arg to decrypt must be chunk or string");
	}

	VALUE vOutput = vInput.Clone();
	p->Decrypt((unsigned long*)vOutput.GetChunkData(), vOutput.GetChunkByteCount() >> 3);
	return vOutput;
}

CORE_FUNC(Blowfish_EncryptVALUE)
{
	REQUIRE_ARGS(1);
	ASSERT_SELF(CRYPTO_MODULE_NAME);

	CBlowfish* p = (CBlowfish*)self->Get(_T("#pBlowfish")).Void();
	if (p == NULL) {
		ERR("Blowfish class has not been instantiated.");
	}

	VALUE vOutput = args[0].SerializeOutMemory();
	if (vOutput.Type() != stChunk) {
		ERR("Failed to serialize value for encrypting");
	}
	p->Encrypt((unsigned long*)vOutput.GetChunkData(), vOutput.GetChunkByteCount() >> 3);
	return vOutput;
}

CORE_FUNC(Blowfish_DecryptVALUE)
{
	REQUIRE_ARGS(1);
	ASSERT_SELF(CRYPTO_MODULE_NAME);

	CBlowfish* p = (CBlowfish*)self->Get(_T("#pBlowfish")).Void();
	if (p == NULL) {
		ERR("Blowfish class has not been instantiated.");
	}

	if (args[0].Type() != stChunk) {
		ERR("Arg to encrypt must be chunk");
	}
	VALUE vInput = args[0].Clone();
	p->Decrypt((unsigned long*)vInput.GetChunkData(), vInput.GetChunkByteCount() >> 3);
	VALUE vOutput;
	if (!vOutput.SerializeInMemory(vInput)) {
		ERR("Failed to re-serialize encrypted value");
	}
	return vOutput;
}

static VALUE InitBlowfishClass()
{
	VALUE v = AllocClass(_T("blowfish"), VALUE());
	v.Set(_T("__init__"), AllocFunction(Blowfish__Init__, CRYPTO_MODULE_NAME));
	v.Set(_T("encrypt"), AllocFunction(Blowfish_Encrypt, CRYPTO_MODULE_NAME));
	v.Set(_T("decrypt"), AllocFunction(Blowfish_Decrypt, CRYPTO_MODULE_NAME));
	v.Set(_T("encryptval"), AllocFunction(Blowfish_EncryptVALUE, CRYPTO_MODULE_NAME));
	v.Set(_T("decryptval"), AllocFunction(Blowfish_DecryptVALUE, CRYPTO_MODULE_NAME));
	return v;
}

/////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////

#include "_bigint.h"

CORE_FUNC(DiffieHellman__Init__)
{
	// diffieHellman(bytesForRand, g, prime)
	//		int: bitsForRand
	//		hex string: g, prime
	//

	if ((argCount != 3) && (argCount != 4)) {
		TERR(lpzWrongNumArgs);
	}

	bigint secret;
	secret.Rand(args[0].Integer());

	self->Set(_T("secret"), AllocChunk(args[0].Integer(), secret.GetBytes()));
	self->Set(_T("g"), AllocChunk(args[1].String(), enc_hex, it_ui1));
	self->Set(_T("prime"), AllocChunk(args[2].String(), enc_hex, it_ui1));
	if (argCount == 4) {
		if (args[3].Type() == stString) {
			self->Set(_T("password"), AllocChunk(args[3].String(), enc_hex, it_ui1));
		}
		else if (args[3].Type() == stChunk) {
			self->Set(_T("password"), args[3].Clone());
		}
	}
	return AllocNil();
}

CORE_FUNC(DiffieHellman_Phase1)
{
	// dh.phase1()
	REQUIRE_ARGS(0);
	ASSERT_SELF(CRYPTO_MODULE_NAME);

	VALUE vSecret = self->Get(_T("secret"));
	VALUE vG = self->Get(_T("g"));
	VALUE vPrime = self->Get(_T("prime"));
	VALUE vPassword = self->Get(_T("password"));

	bigint Y;
	bigint secret = bigint(vSecret.GetChunkData(), vSecret.GetChunkByteCount(), true);
	bigint prime = bigint(vPrime.GetChunkData(), vPrime.GetChunkByteCount(), true);
	if (vPassword.Type() == stChunk) {
		// Do:  Y = power(password, 2 * secret) mod p
		secret.LeftShiftIP(); // secret = secret * 2
		bigint password = bigint(vPassword.GetChunkData(), vPassword.GetChunkByteCount(), true);
		Y = bigint::ModularExp(password, secret, prime).ToString();
	}
	else {
		// Standard Diffie-Hellman
		// Do:  Y = power(g, secret) mod p
		bigint g = bigint(vG.GetChunkData(), vG.GetChunkByteCount(), true);
		Y = bigint::ModularExp(g, secret, prime).ToString();
	}
	return Y.ToString();
}

CORE_FUNC(DiffieHellman_Phase2)
{
	// dh.phase2(Y)
	REQUIRE_ARGS(1);
	ASSERT_SELF(CRYPTO_MODULE_NAME);

	VALUE vSecret = self->Get(_T("secret"));
	VALUE vPrime = self->Get(_T("prime"));
	VALUE vPassword = self->Get(_T("password"));

	bigint Y = bigint(args[0].String());
	bigint secret = bigint(vSecret.GetChunkData(), vSecret.GetChunkByteCount(), true);
	bigint prime = bigint(vPrime.GetChunkData(), vPrime.GetChunkByteCount(), true);
	if (vPassword.Type() == stChunk) {
		// Do:  Y = power(Y, 2 * secret) mod p
		secret.LeftShiftIP();
	}
	else {
		// Standard Diffie-Hellman
	}
	return bigint::ModularExp(Y, secret, prime).ToString();
}

static VALUE InitDiffieHellmanClass()
{
	VALUE v = AllocClass(_T("diffiehellman"), VALUE());
	v.Set(_T("__init__"), AllocFunction(DiffieHellman__Init__, CRYPTO_MODULE_NAME));
	v.Set(_T("phase1"), AllocFunction(DiffieHellman_Phase1, CRYPTO_MODULE_NAME));
	v.Set(_T("phase2"), AllocFunction(DiffieHellman_Phase2, CRYPTO_MODULE_NAME));
	return v;
}

/////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////

#include "MD5.h"

static void freeMD5(void* pVoid)
{
	if (pVoid) delete (MD5_CTX*)pVoid;
}

CORE_FUNC(MD5__Init__)
{
	REQUIRE_ARGS(0);
	ASSERT_SELF(CRYPTO_MODULE_NAME);

	MD5_CTX* p = new MD5_CTX;
	MD5Init(p);
	self->Set(_T("#pMD5"), AllocVoid(p, freeMD5, _T("md5")));
	return AllocNil();
}

CORE_FUNC(MD5_Update)
{
	REQUIRE_ARGS(1);
	ASSERT_SELF(CRYPTO_MODULE_NAME);

	MD5_CTX* p = (MD5_CTX*)self->Get(_T("#pMD5")).Void();
	if (p == NULL) {
		ERR("MD5 class has not been instantiated.");
	}

	if (args[0].Type() == stChunk) {
		MD5Update(p, args[0].GetChunkData(), args[0].GetChunkByteCount());
	}
	else {
#ifdef UNDER_CE
		CString text = args[0].String();
		int length = text.GetLength();
		BYTE *buf = new BYTE[length];
		for (int idx = 0; idx < length; ++idx) {
			buf[idx] = (BYTE)(text[idx]);
		}
		MD5Update(p, buf, length);
		delete[] buf;
#else
		CString text = args[0].String();
		MD5Update(p, (BYTE*)(LPCTSTR)text, text.GetLength());
#endif
	}
	return AllocNil();
}

CORE_FUNC(MD5_Digest)
{
	REQUIRE_ARGS(0);
	ASSERT_SELF(CRYPTO_MODULE_NAME);

	MD5_CTX* p = (MD5_CTX*)self->Get(_T("#pMD5")).Void();
	if (p == NULL) {
		ERR("MD5 class has not been instantiated.");
	}

	unsigned char digest[16];
	MD5Final (digest, p);
	return AllocChunk(16, digest);
}

static VALUE InitMD5Class()
{
	VALUE v = AllocClass(_T("md5"), VALUE());
	v.Set(_T("__init__"), AllocFunction(MD5__Init__, CRYPTO_MODULE_NAME));
	v.Set(_T("update"), AllocFunction(MD5_Update, CRYPTO_MODULE_NAME));
	v.Set(_T("digest"), AllocFunction(MD5_Digest, CRYPTO_MODULE_NAME));
	return v;
}

/////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////

#include "sha.h"

static void freeSHA(void* pVoid)
{
	if (pVoid) delete (SHAobject*)pVoid;
}

CORE_FUNC(SHA__Init__)
{
	REQUIRE_ARGS(0);
	ASSERT_SELF(CRYPTO_MODULE_NAME);

	SHAobject* p = new SHAobject;
	SHAinit(p);
	self->Set(_T("#pSHA"), AllocVoid(p, freeSHA, _T("sha1")));
	return AllocNil();
}

CORE_FUNC(SHA_Update)
{
	REQUIRE_ARGS(1);
	ASSERT_SELF(CRYPTO_MODULE_NAME);

	SHAobject* p = (SHAobject*)self->Get(_T("#pSHA")).Void();
	if (p == NULL) {
		ERR("SHA class has not been instantiated.");
	}

	if (args[0].Type() == stChunk) {
		SHAupdate(p, args[0].GetChunkData(), args[0].GetChunkByteCount());
	}
	else {
#ifdef UNDER_CE
		CString text = args[0].String();
		int length = text.GetLength();
		BYTE *buf = new BYTE[length];
		for (int idx = 0; idx < length; ++idx) {
			buf[idx] = (BYTE)(text[idx]);
		}
		SHAupdate(p, buf, length);
		delete[] buf;
#else
		CString text = args[0].String();
		SHAupdate(p, (BYTE*)(LPCTSTR)text, text.GetLength());
#endif
	}
	return AllocNil();
}

CORE_FUNC(SHA_Digest)
{
	REQUIRE_ARGS(0);
	ASSERT_SELF(CRYPTO_MODULE_NAME);

	SHAobject* p = (SHAobject*)self->Get(_T("#pSHA")).Void();
	if (p == NULL) {
		ERR("MD5 class has not been instantiated.");
	}

	SHAFinal(p);
	unsigned char digest[SHS_DIGESTSIZE];
	for(int i = 0, j = 0; i<SHS_DIGESTSIZE; i+=4, j++) {
		digest[i+3]=(BYTE)((p->digest[j]      ) & 255);
		digest[i+2]=(BYTE)((p->digest[j] >>  8) & 255);
		digest[i+1]=(BYTE)((p->digest[j] >> 16) & 255);
		digest[i  ]=(BYTE)((p->digest[j] >> 24) & 255);
	}
	return AllocChunk(SHS_DIGESTSIZE, digest);
}

static VALUE InitSHAClass()
{
	VALUE v = AllocClass(_T("sha1"), VALUE());
	v.Set(_T("__init__"), AllocFunction(SHA__Init__, CRYPTO_MODULE_NAME));
	v.Set(_T("update"), AllocFunction(SHA_Update, CRYPTO_MODULE_NAME));
	v.Set(_T("digest"), AllocFunction(SHA_Digest, CRYPTO_MODULE_NAME));
	return v;
}

/////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////

VALUE InitCryptoModule()
{
	VALUE module = AllocFrame();
	module.Set(_T("__module__"), CRYPTO_MODULE_NAME);

	vBlowfishClassBase = InitBlowfishClass();
	module.Set(_T("blowfish"), vBlowfishClassBase);
	vDiffieHellmanClassBase = InitDiffieHellmanClass();
	module.Set(_T("diffiehellman"), vDiffieHellmanClassBase);
	vMD5ClassBase = InitMD5Class();
	module.Set(_T("md5"), vMD5ClassBase);
	vSHAClassBase = InitSHAClass();
	module.Set(_T("sha1"), vSHAClassBase);

	return module;
}

#endif // #ifdef CRYPTO_SUPPORT
